library(data.table)
library(gridExtra)
library(grid)
library(dplyr)
library(tidyr)
library(ggplot2)
library(stringr)
library(readr)

## Paper or presentation
paper <- TRUE

## data directory
dir.project <- './'
dir.data <- paste0(dir.project, 'data/')
dir.raw <- paste0(dir.data, 'raw/')
dir.generated <- paste0(dir.data, 'generated/')

## figures directory for output
dir.fig.paper <-
  paste0(dir.project, 'doc/paper/figures/compensation/')

source(paste0(dir.project, 'src/r/figures.R'))

fileName <- paste0(dir.project, 'calibration_version.txt')
calibration_version <-
  read_file(fileName) %>% str_replace_all("[\r\n]" , "")

#############
## FIGURES ##
#############

## load data
df.status_quo <-
  fread(
    paste0(
      dir.generated,
      'matlab/simulation/20191107_nochoice_simulation.csv'
    )
  ) %>% mutate(version = 'status_quo') %>%
  mutate(welfare_summed = 
           welfare_part_M +
           welfare_part_R + 
           welfare_part_C + 
           welfare_non_part) %>%
  mutate(welfare_comparison = welfare_summed - welfare)

df.adjust_robottax <-
  fread(
    paste0(
      dir.generated,
      'matlab/OptimTax/hsv/20191107_nochoice_adjust_robot_tax.csv'
    )
  ) %>% mutate(version = 'adjust_robottax') %>%
  mutate(welfare_summed = 
           welfare_part_M +
           welfare_part_R + 
           welfare_part_C + 
           welfare_non_part) %>%
  mutate(welfare_comparison = welfare_summed - welfare)

df.optimum <-
    fread(paste0(dir.generated, 'matlab/OptimTax/gpops/', 'optimum_vars.csv')) %>%
  mutate(version = 'optimum', pricedrop = drop) %>%
  mutate(welfare_summed = 
           welfare_M +
           welfare_R + 
           welfare_C + 
           welfare_nonpart) %>%
  mutate(welfare_comparison = welfare_summed - welfare) %>%
  rename(welfare_part_M = welfare_M,
         welfare_part_R = welfare_R,
         welfare_part_C = welfare_C,
         welfare_non_part = welfare_nonpart)

df.combined <- rbind(df.status_quo %>% select(version, pricedrop, starts_with('welfare'), part, part_M, part_R, part_C, starts_with('pop'), starts_with('emp')),
                     df.adjust_robottax %>% select(version, pricedrop, starts_with('welfare'), part, part_M, part_R, part_C, starts_with('pop'), starts_with('emp')),
                     df.optimum %>% select(version, pricedrop, starts_with('welfare'), part, part_M, part_R, part_C, starts_with('pop'), starts_with('emp')), fill=TRUE)


ggplot(df.combined, aes(x=pricedrop, y=welfare)) + geom_line(aes(color=version))

# compute various shares
df.combined <- df.combined %>%
  rename(
    welfare_total_part_M = welfare_part_M,
    welfare_total_part_R = welfare_part_R,
    welfare_total_part_C = welfare_part_C
  ) %>%
  mutate(pricedrop = pricedrop*100) %>%
  mutate(non_part_M = pop_M - part_M,
         non_part_R = pop_R - part_R,
         non_part_C = pop_C - part_C,
         non_part = 1 - part,
         non_emp_M = non_part_M / non_part,
         non_emp_R = non_part_R / non_part,
         non_emp_C = non_part_C / non_part,
         welfare_total_nonpart_M = non_emp_M * welfare_non_part,
         welfare_total_nonpart_R = non_emp_R * welfare_non_part,
         welfare_total_nonpart_C = non_emp_C * welfare_non_part,
         welfare_percap_part_M = welfare_total_part_M / part_M,
         welfare_percap_part_R = welfare_total_part_R / part_R,
         welfare_percap_part_C = welfare_total_part_C / part_C,
         welfare_percap_nonpart_M = welfare_total_nonpart_M / non_part_M,
         welfare_percap_nonpart_R = welfare_total_nonpart_R / non_part_R,
         welfare_percap_nonpart_C = welfare_total_nonpart_C / non_part_C,
         welfare_total_all_M = welfare_total_part_M + welfare_total_nonpart_M,
         welfare_total_all_R = welfare_total_part_R + welfare_total_nonpart_R,
         welfare_total_all_C = welfare_total_part_C + welfare_total_nonpart_C,
         welfare_percap_all_M = welfare_total_all_M / pop_M,
         welfare_percap_all_R = welfare_total_all_R / pop_R,
         welfare_percap_all_C = welfare_total_all_C / pop_C
         ) %>%
    rename(
      part_total_all_M = part_M,
      part_total_all_R = part_R,
      part_total_all_C = part_C) %>%
  mutate(
    part_rate_all_M = part_total_all_M / pop_M * 100,
    part_rate_all_R = part_total_all_R / pop_R * 100,
    part_rate_all_C = part_total_all_C / pop_C * 100,
    ) %>%
  select(-welfare, -welfare_summed, -welfare_comparison, -welfare_part, -welfare_non_part) %>%
  select(pricedrop, version, starts_with('welfare'), starts_with('part_'))


# transform to long

df.combined.long <- df.combined %>% pivot_longer(-c(pricedrop,version), names_sep='_', names_to=c('measure', 'type', 'lm_status', 'occupation'))

df.combined.long$occupation <- factor(df.combined.long$occupation, levels=c('M','R','C'), labels=c('manual', 'routine', 'cognitive'))
df.combined.long$lm_status <- factor(df.combined.long$lm_status, levels=c('part','nonpart','all'), labels=c('participants', 'non-participants', 'all'))
df.combined.long$type <- factor(df.combined.long$type, levels=c('total','percap','rate'))
df.combined.long$version <- factor(df.combined.long$version, levels=c('status_quo','adjust_robottax', 'optimum'),
                                   labels=c('status-quo', 'status-quo + opt. robot tax', 'optimum'))

p.welfare <- ggplot(df.combined.long %>% filter(type=='percap', lm_status!='all', measure=='welfare'), aes(x=pricedrop, y=value)) +
    geom_line(aes(linetype=version)) + facet_grid(cols = vars(occupation), rows=vars(lm_status)) +
  theme_bw() + theme.serif + ylab('') + xlab('pct. drop in robot price') + theme(legend.position = 'top', legend.title = element_blank()) +
  theme(axis.text.y = element_blank())

p.part <- ggplot(df.combined.long %>% filter(lm_status=='all', type=='rate', measure=='part'), aes(x=pricedrop, y=value)) +
  geom_line(aes(linetype=version)) + facet_grid(cols = vars(occupation)) +
  theme_bw() + theme.serif + ylab('') + xlab('pct. drop in robot price') + theme(legend.position = 'top', legend.title = element_blank())

ggsave(
  paste0(dir.fig, "plot_gains_losses_welfare.pdf"),
  plot = p.welfare,
  width = 12,
  height = 10,
  units = "cm"
)

ggsave(
  paste0(dir.fig, "plot_gains_losses_participation.pdf"),
  plot = p.part,
  width = 12,
  height = 6,
  units = "cm"
)

## additional not-used plots

ggplot(df.combined.long %>% filter(type=='total', lm_status!='all', measure=='welfare'), aes(x=pricedrop, y=value)) +
  geom_line(aes(linetype=version)) + facet_grid(cols = vars(occupation), rows=vars(lm_status)) +
  theme_bw() + theme.serif + ylab('') + xlab('pct. drop in robot price')

ggplot(df.combined.long %>% filter(lm_status=='all', type=='percap', measure=='welfare'), aes(x=pricedrop, y=value)) +
  geom_line(aes(linetype=version)) + facet_grid(cols = vars(occupation)) +
  theme_bw() + theme.serif + ylab('') + xlab('pct. drop in robot price')
